package org.groupg.practice;

import org.jgrapht.Graph;
//import org.jgrapht.graph.DefaultDirectedAcyclicGraph;
import org.jgrapht.graph.DefaultDirectedGraph;
import org.jgrapht.graph.DefaultEdge;
import org.jgrapht.traverse.TopologicalOrderIterator;

import java.util.*;

public class DagLayerExample {

    public static void main(String[] args) {
        // 创建一个有向无环图
        Graph<String, DefaultEdge> dag = new DefaultDirectedGraph<>(DefaultEdge.class);

        // 添加顶点
        dag.addVertex("A");
        dag.addVertex("B");
        dag.addVertex("C");
        dag.addVertex("D");
        dag.addVertex("E");
        dag.addVertex("F");
        dag.addVertex("G");

        // 添加边
        dag.addEdge("A", "B");
        dag.addEdge("G", "B");
        dag.addEdge("B", "C");
        dag.addEdge("C", "D");
        dag.addEdge("D", "E");
        dag.addEdge("E", "F");
        dag.addEdge("E", "F");

        // 执行拓扑排序
        TopologicalOrderIterator<String, DefaultEdge> topoIter = new TopologicalOrderIterator<>(dag);

        // 使用一个映射来存储每一层的顶点
        Map<Integer, List<String>> layers = new LinkedHashMap<>(); // 保持插入顺序
        int layerIndex = 0;

        // 使用一个队列来进行广度优先搜索(BFS)式的分层
        Queue<String> currentLayer = new LinkedList<>();
        if (topoIter.hasNext()) {
            currentLayer.add(topoIter.next());
        }

        while (!currentLayer.isEmpty()) {
            int size = currentLayer.size();
            List<String> currentLayerList = new ArrayList<>();

            for (int i = 0; i < size; i++) {
                String vertex = currentLayer.poll();
                currentLayerList.add(vertex);

                // 找到当前顶点的所有后继顶点并加入下一层
                for (String successor : dag.outgoingEdgesOf(vertex).stream()
                        .map(dag::getEdgeTarget)
                        .toList()) {
                    // 只有当后继顶点还没被处理过时才加入下一层
                    if (!layers.values().stream().flatMap(List::stream).toList().contains(successor) &&
                        !currentLayer.contains(successor)) {
                        if (!currentLayer.containsAll(dag.outgoingEdgesOf(successor).stream()
                                .map(dag::getEdgeTarget)
                                .toList())) { // 简单检查避免同一层重复加入
                            ((LinkedList<String>) ((LinkedHashMap<Integer, List<String>>) layers).getOrDefault(layerIndex + 1, new LinkedList<>())).add(successor);
                        }
                    }
                }
            }

            layers.put(layerIndex, currentLayerList);
            layerIndex++;

            // 下一层成为当前层继续处理
            if (!layers.containsKey(layerIndex)) {
                currentLayer = new LinkedList<>();
            } else {
                currentLayer = (Queue<String>) layers.get(layerIndex);
                layers.remove(layerIndex); // 从map中移除，因为我们要用queue来处理它
                layerIndex--; // 因为我们马上又要处理这一层，所以index不变实际是相当于没加
            }

            // 如果topoIter还有元素且currentLayer为空，说明还有顶点没加入任何层（应该是下一层的开始）
            if (topoIter.hasNext() && currentLayer.isEmpty()) {
                currentLayer.add(topoIter.next());
            }
        }

        // 输出每一层的顶点
        for (Map.Entry<Integer, List<String>> entry : layers.entrySet()) {
            System.out.println("Layer " + entry.getKey() + ": " + entry.getValue());
        }
    }
}